# 数据集准备
x = torch.tensor([-2.7920, -2.2356, -1.8285, -1.7755, -1.7408, -1.0992,
                  -0.9828, -0.9527, -0.2743, -0.2406, -0.0918, 0.1785,
                  0.2922, 0.3702, 0.3957, 0.4205, 0.5041, 0.5499,
                  0.6154, 0.9958, 1.2607, 1.2784, 1.4159, 1.9559])
x_input = x.reshape(-1, 1)
y_target = 0.5 * x ** 1 + 3. * x ** 2 + 0.4 * x ** 3  # (24,1)
y_true = y_target.reshape(-1, 1)

# 构建模型
class poly_model(nn.Module):
    def __init__(self):
        super(poly_model, self).__init__()
        self.hidden1 = torch.nn.Linear(1, 20)
        self.hidden2 = torch.nn.Linear(20, 20)
        self.hidden3 = torch.nn.Linear(20, 1)
    def forward(self, x):
        x = torch.relu(self.hidden1(x))  # 用ReLU激活函数
        x = torch.relu(self.hidden2(x))
        out = self.hidden3(x)
        return out
model = poly_model()
criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001)
for epoch in range(3000):
    out = model(x_input)
    loss = criterion(out, y_true)
    print_loss = loss.item()
    # backward
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print('epoch is {} loss is {}'.format(epoch, print_loss))
# 预测结果
model.eval()
plt.rcParams["font.sans-serif"] = ["SimHei"]  # 设置字体
plt.rcParams["axes.unicode_minus"] = False  # 正常显示负号
fig, ax = plt.subplots(1, 1)
ax.plot(x, y_target.squeeze().numpy(), 'ob')
predict = model(x_input)
predict = predict.data.numpy()
ax.plot(x, predict, label='多层神经网络拟合')
plt.xlabel('X')
plt.ylabel('y')
ax.legend(loc='upper center', frameon=False)
plt.show() 
